import cv2
import numpy as np
from utils.img import load_gray
from utils.logger import Logger
from datetime import datetime
from utils.folder import ensure_empty_directory

class DeviceBase:
    def __init__(self):
        self.logger = Logger()
        self.screen = None
        self.screenshot()
        self.screen_width = self.screen.shape[1]
        self.screen_height = self.screen.shape[0]
        self.screen_center = (self.screen_width // 2, self.screen_height // 2)

        self.record_dir  = "records/"
        ensure_empty_directory(self.record_dir)

    def tap(self, x:int, y:int):
        raise NotImplementedError()
    
    def back(self):
        raise NotImplementedError()
        
    def long_press(self, x:int, y:int, seconds:int):
        raise NotImplementedError()
    
    def swipe_up(self, ratio:float, x:int):
        raise NotImplementedError()
    
    def swipe_down(self, ratio:float, x:int):
        raise NotImplementedError()
    
    def swipe_left(self, ratio:float, y:int):
        raise NotImplementedError()
    
    def swipe_right(self, ratio:float, y:int):
        raise NotImplementedError()

    def screenshot(self):
        '''
        更新对象的screen属性
        '''
        raise NotImplementedError()
    
    def save_screen(self, path:str="screen.png"):
        cv2.imwrite(path, self.screen)
        print(f'save screen to {path}')
    
    def find_image(self, template, threshold=0.8, x_start:int=0, x_end:int=None, y_start:int=0, y_end:int=None, ratio:float=0.25) -> list[tuple[int, int]]:
        if x_end is None:
            x_end = self.screen_width
        if y_end is None:
            y_end = self.screen_height

        self.screenshot()
        ROI = self.screen[y_start:y_end, x_start:x_end]
        screen = cv2.cvtColor(ROI, cv2.COLOR_RGB2GRAY)
        screen = cv2.resize(screen, fx=ratio, fy=ratio)

        if isinstance(template, str):
            template = load_gray(template)
        elif len(template.shape) > 2:
            template = cv2.cvtColor(template, cv2.COLOR_RGB2GRAY)
        template = cv2.resize(template, fx=ratio, fy=ratio)

        res = cv2.matchTemplate(screen, template, cv2.TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)
        
        ans = []
        for pt in zip(*loc[::-1]):
            if all((pt[0] - x)**2 + (pt[1] - y)**2 > min(template.shape) for (x, y) in ans):
                ans.append(pt)
        
        for i in range(len(ans)):
            ans[i] = (
                int((ans[i][0] + template.shape[1] / 2) / ratio + x_start), 
                int((ans[i][1] + template.shape[0] / 2) / ratio + y_start)
                )
        return ans
    
    def find_image_SIFT(self, template, threshold=0.8, x_start:int=0, x_end:int=None, y_start:int=0, y_end:int=None, ratio:float=0.25) -> list[tuple[int, int]]:
        if x_end is None:
            x_end = self.screen_width
        if y_end is None:
            y_end = self.screen_height
        
        self.screenshot()
        ROI = self.screen[y_start:y_end, x_start:x_end]
        screen = cv2.cvtColor(ROI, cv2.COLOR_RGB2GRAY)
        screen = cv2.resize(screen, fx=ratio, fy=ratio)
        
        if isinstance(template, str):
            template = load_gray(template)
        elif len(template.shape) > 2:
            template = cv2.cvtColor(template, cv2.COLOR_RGB2GRAY)
        template = cv2.resize(template, fx=ratio, fy=ratio)

        # 初始化SIFT检测器
        sift = cv2.SIFT_create()

        # 查找模板图像和目标图像的关键点和描述符
        kp1, des1 = sift.detectAndCompute(template, None)
        kp2, des2 = sift.detectAndCompute(screen, None)

        # 创建BFMatcher对象
        bf = cv2.BFMatcher()

        # 使用KNN匹配
        matches = bf.knnMatch(des1, des2, k=2)
        # 应用比率测试来选择好的匹配
        good_matches = []
        for m, n in matches:
            if m.distance < 0.75 * n.distance:
                good_matches.append([
                        int((kp2[m.trainIdx].pt[0] + template.shape[1] / 2) / ratio + x_start), 
                        int((kp2[m.trainIdx].pt[1] + template.shape[0] / 2) / ratio + y_start)
                    ])

        return good_matches
    
    def waiting_screen_changed(self, threshold=0.8, func:callable=None, args:set=set()):
        pre = cv2.cvtColor(self.screen, cv2.COLOR_RGB2GRAY)
        func(*args) if func else None
        sim = 1
        while sim > threshold:
            self.screenshot()
            now = cv2.cvtColor(self.screen, cv2.COLOR_RGB2GRAY)
            # 基于统计计算相似度
            sim = self.simalrity(pre, now)

    def simalrity(self, img1, img2) -> float:
        if len(img1.shape) > 2:
            img1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
        if len(img2.shape) > 2:
            img2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
        return np.sum((img1-img2)==0) / (img1.shape[0] * img1.shape[1])
    
    def tap_image(self, img, all=True, **args) -> int:
        points = self.find_image(img, **args)
        if len(points) == 0:
            return 0
        if not all:
            self.tap(*points[0])
        else:
            for point in points:
                self.tap(*point)
        return len(points)

    def record(self, x, y, op:str, ratio:float=0.25):
        self.screenshot()
        img = cv2.cvtColor(self.screen, cv2.COLOR_RGB2GRAY)
        cv2.line(img, (0, y), (self.screen_width, y), (0, 255, 0), 2)
        cv2.line(img, (x, 0), (x, self.screen_height), (0, 255, 0), 2)        
        file_name = f"{self.record_dir}{datetime.now().strftime("%Y%m%d-%H%M%S")}_{op}_{x}_{y}.png"
        img = cv2.resize(img, fx=ratio, fy=ratio)
        cv2.imwrite(file_name, img)